module deform_test_mod

  ! reference: Nair and Lauritzen,2010. A class of deformational flow test cases for linear transport problems on the sphere. JCP

	use const_mod
	use namelist_mod
	use sphere_geometry_mod
	use mesh_mod
	use state_mod
	use tracer_mod
	use parallel_mod

	implicit none

	private
	public deform_test_init
	public deform_test_set_ic
	public deform_case1_test_set_uv
	public deform_case2_test_set_uv
	public deform_case3_test_set_uv
	public deform_case4_test_set_uv

	real(r8), parameter :: period = 12 * 86400
	real(r8) lon1, lat1, lon2, lat2

contains

	subroutine deform_test_init(case_n)

		integer, intent(in) :: case_n

		select case (case_n)
		case (1)
			lon1 =  pi
			lat1 =  pi / 3.0_r8
			lon2 =  pi
			lat2 = -pi / 3.0_r8
		case (3)
			lon1 = pi * 3.0_r8 / 4.0_r8
			lat1 = 0
			lon2 = pi * 5.0_r8 / 4.0_r8
			lat2 = 0
		case (2, 4)
			lon1 = pi * 5.0_r8 / 6.0_r8
			lat1 = 0
			lon2 = pi * 7.0_r8 / 6.0_r8
			lat2 = 0
		end select

	end subroutine deform_test_init

	subroutine deform_test_set_ic(mesh, tracer)

		type(mesh_type), intent(in) :: mesh
		type(tracer_type), intent(inout) :: tracer
		integer i, j, itracer
    real(r8) lon, lat, r, r1, r2, qmax, qmin, c
    real(r8) x(3), x1(3), x2(3)

    call cartesian_transform(lon1, lat1, x1(1), x1(2), x1(3)); x1 = x1 / radius
    call cartesian_transform(lon2, lat2, x2(1), x2(2), x2(3)); x2 = x2 / radius
    if (ntracers /= 4) then
    	print*, 'deform_test case need 4 tracers!'
    	stop
    end if

    ! Background
    tracer%q(:,:,:,1) = 1.d0
    ! Cosine hills
    qmax = 1.0_r8; qmin = 0.1_r8; c = 0.9_r8; r = radius * 0.5_r8
    do j = mesh%full_lat_ibeg, mesh%full_lat_iend
    	lat = mesh%full_lat(j)
    	do i = mesh%full_lon_ibeg, mesh%full_lon_iend
    		lon = mesh%full_lon(i)
    		r1 = calc_distance(lon1, lat1, lon, lat)
    		r2 = calc_distance(lon2, lat2, lon, lat)
    		if (r1 < r) then
    			tracer%q(i,j,1,2) = qmin + c * qmax * 0.5_r8 * (1 + cos(pi * r1 / r))
    		else if (r2 < r) then
    			tracer%q(i,j,1,2) = qmin + c * qmax * 0.5_r8 * (1 + cos(pi * r2 / r))
    		else
    			tracer%q(i,j,1,2) = qmin
    		end if
    	end do
    end do
    ! Slotted cylinders
    qmax = 1.0_r8; qmin = 0.1_r8; r = radius * 0.5_r8
    do j = mesh%full_lat_ibeg, mesh%full_lat_iend
    	lat = mesh%full_lat(j)
    	do i = mesh%full_lon_ibeg, mesh%full_lon_iend
    		lon = mesh%full_lon(i)
    		r1 = calc_distance(lon1, lat1, lon, lat)
    		r2 = calc_distance(lon2, lat2, lon, lat)
    		if ((r1 <= r .and. abs(lon - lon1) >= r / radius / 6.0_r8) .or. &
    		    (r2 <= r .and. abs(lon - lon2) >= r / radius / 6.0_r8)) then
    			tracer%q(i,j,1,3) = qmax
    		else if (r1 <= r .and. abs(lon - lon1) < r / radius / 6.0_r8 .and. &
    			 lat - lat1 < -5.0_r8 / 12.0_r8 * (r / radius)) then
    			tracer%q(i,j,1,3) = qmax
    		else if (r2 <= r .and. abs(lon - lon2) < r / radius / 6.0_r8 .and. &
    			 lat - lat2 > 5.0_r8 / 12.0_r8 * (r / radius)) then
    			tracer%q(i,j,1,3) = qmax
    		else
    			tracer%q(i,j,1,3) = qmin
    		end if
    	end do
    end do 
    ! Gaussin hills
    qmax = 0.95_r8; c = 5.0_r8
    do j = mesh%full_lat_ibeg, mesh%full_lat_iend
    	lat = mesh%full_lat(j)
    	do i = mesh%full_lon_ibeg, mesh%full_lon_iend
    		lon = mesh%full_lon(i)
    		call cartesian_transform(lon, lat, x(1), x(2), x(3))
    		x = x / radius
    		tracer%q(i,j,1,4) = qmax * (exp(-c * dot_product(x - x1, x - x1)) +&
    		                            exp(-c * dot_product(x - x2, x - x2)))
    	end do
    end do 

    do itracer = 1, size(tracer%q, 4)
    	call fill_zonal_halo_cell(tracer%q(:,:,:,itracer), all_halo=.true.)
    	call fill_merid_halo_cell(tracer%q(:,:,:,itracer), all_halo=.true.)
    end do

	end subroutine deform_test_set_ic

	subroutine deform_case1_test_set_uv(mesh, state, time_in_seconds)

		type(mesh_type), intent(in) :: mesh
		type(state_type), intent(inout) :: state
		real(8), intent(in) :: time_in_seconds

		integer i, j
		real(r8) psi0, lon, lat, cos_t, k0

		associate (u => state%u, v => state%v, psi => state%psi)
		psi0 = 10.0_r8  * radius / period
		cos_t = cos(pi * time_in_seconds / period)
		k0 = 10.0_r8 * radius / period

		! do j = mesh%half_lat_ibeg, mesh%half_lat_iend
		! 	lat = mesh%half_lat(j)
		! 	do i = mesh%half_lon_ibeg, mesh%half_lon_iend
		! 		lon = mesh%half_lon(i)
		! 		psi(i,j) = psi0 * sin(lon / 2)**2 * cos(lat)**2 * cos_t
		! 	end do
		! end do

		! call fill_zonal_halo_lat(psi, all_halo=.true.)

		! do j = mesh%full_lat_ibeg + 1, mesh%full_lat_iend - 1
		! 	do i = mesh%half_lon_ibeg, mesh%half_lon_iend
		! 		u(i,j,1) = (psi(i,j) - psi(i,j-1)) / mesh%le_lon(j) * radius
		! 	end do
		! end do
  !   call fill_zonal_halo_lon(u, all_halo=.true.)
		
		! do j = mesh%half_lat_ibeg, mesh%half_lat_iend
		! 	do i = mesh%full_lon_ibeg, mesh%full_lon_iend
		! 		v(i,j,1) = - (psi(i,j) - psi(i-1,j)) / mesh%le_lat(j) * radius
		! 	end do
		! end do

		do j = mesh%full_lat_ibeg + 1, mesh%full_lat_iend - 1
			lat = mesh%full_lat(j)
			do i = mesh%half_lon_ibeg, mesh%half_lon_iend
				lon = mesh%half_lon(i)
				u(i,j,1) = k0 * sin(lon / 2.0_r8)**2 * sin(2 * lat) * cos_t
			end do
		end do
		call fill_zonal_halo_lon(u, all_halo=.true.)

		do j = mesh%half_lat_ibeg, mesh%half_lat_iend
			lat = mesh%half_lat(j)
			do i = mesh%full_lon_ibeg, mesh%full_lon_iend
				lon = mesh%full_lon(i)
				v(i,j,1) = k0 / 2.0_r8 * sin(lon) * cos(lat) * cos_t
			end do
		end do
		end associate

	end subroutine deform_case1_test_set_uv
	
	subroutine deform_case2_test_set_uv(mesh, state, time_in_seconds)

		type(mesh_type), intent(in) :: mesh
		type(state_type), intent(inout) :: state
		real(8), intent(in) :: time_in_seconds

		integer i, j
		real(r8) lon, lat, k0, cos_t

		associate (u => state%u, v => state%v, psi => state%psi)
		k0 = 10.0_r8 * radius / period
		cos_t = cos(pi * time_in_seconds / period)

		! do j = mesh%half_lat_ibeg, mesh%half_lat_iend
		! 	lat = mesh%half_lat(j)
		! 	do i = mesh%half_lon_ibeg, mesh%half_lon_iend
		! 		lon = mesh%half_lon(i)
		! 		psi(i,j) = k0 * sin(lon)**2 * cos(lat)**2 * cos_t
		! 	end do
		! end do

		! call fill_zonal_halo_lat(psi, all_halo=.true.)

		! do j = mesh%full_lat_ibeg + 1, mesh%full_lat_iend - 1
		! 	do i = mesh%half_lon_ibeg, mesh%half_lon_iend
		! 		u(i,j,1) = (psi(i,j) - psi(i,j-1)) / mesh%le_lon(j) * radius
		! 	end do
		! end do
  !   call fill_zonal_halo_lon(u, all_halo=.true.)
		
		! do j = mesh%half_lat_ibeg, mesh%half_lat_iend
		! 	do i = mesh%full_lon_ibeg, mesh%full_lon_iend
		! 		v(i,j,1) = - (psi(i,j) - psi(i-1,j)) / mesh%le_lat(j) * radius
		! 	end do
		! end do
		do j = mesh%full_lat_ibeg + 1, mesh%full_lat_iend - 1
			lat = mesh%full_lat(j)
			do i = mesh%half_lon_ibeg, mesh%half_lon_iend
				lon = mesh%half_lon(i)
				u(i,j,1) = k0 * sin(lon)**2 * sin(2 * lat) * cos_t
			end do
		end do

		call fill_zonal_halo_lon(u, all_halo=.true.)

		do j = mesh%half_lat_ibeg, mesh%half_lat_iend
			lat = mesh%half_lat(j)
			do i = mesh%full_lon_ibeg, mesh%full_lon_iend
				lon = mesh%full_lon(i)
				v(i,j,1) = k0 * sin(2 * lon) * cos(lat) * cos_t
			end do
		end do

		end associate

	end subroutine deform_case2_test_set_uv

	subroutine deform_case3_test_set_uv(mesh, state, time_in_seconds)

		type(mesh_type), intent(in) :: mesh
		type(state_type), intent(inout) :: state
		real(8), intent(in) :: time_in_seconds

		integer i, j
		real(r8) lon, lat, k0, cos_t

		associate (u => state%u, v => state%v)
		k0 = 5.0_r8 * radius / period
		cos_t = cos(pi * time_in_seconds / period)
		do j = mesh%full_lat_ibeg + 1, mesh%full_lat_iend - 1
			lat = mesh%full_lat(j)
			do i = mesh%half_lon_ibeg, mesh%half_lon_iend
				lon = mesh%half_lon(i)
				u(i,j,1) = -k0 * sin(lon / 2)**2 * sin(2 * lat) * cos(lat)**2 * cos_t
			end do 
		end do

		call fill_zonal_halo_lon(u, all_halo=.true.)

		do j = mesh%half_lat_ibeg, mesh%half_lat_iend
			lat = mesh%half_lat(j)
			do i = mesh%full_lon_ibeg, mesh%full_lon_iend
				lon = mesh%full_lon(i)
				v(i,j,1) = k0 / 2.0_r8 * sin(lon) * cos(lat)**3 * cos_t
			end do
		end do 

		end associate
			
	end subroutine deform_case3_test_set_uv

	subroutine deform_case4_test_set_uv(mesh, state, time_in_seconds)

		type(mesh_type), intent(in) :: mesh
		type(state_type), intent(inout) :: state
		real(8), intent(in) :: time_in_seconds

		integer i, j
		real(r8) lon, lat, k0, c1, c2, cos_t

		associate (u => state%u, v => state%v, psi => state%psi)
		k0 = 10._r8 * radius / period
		c1 = pi2 * time_in_seconds / period
		c2 = pi2 * radius / period
		cos_t = cos(pi * time_in_seconds / period)

		! do j = mesh%half_lat_ibeg, mesh%half_lat_iend
		! 	lat = mesh%half_lat(j)
		! 	do i = mesh%half_lon_ibeg, mesh%half_lon_iend
		! 		lon = mesh%half_lon(i) - c1
		! 		psi(i,j) = -k0 * sin(lon)**2 * cos(lat)**2 * cos_t + pi2 * sin(lat) / period * radius
		! 	end do
		! end do
		! call fill_zonal_halo_lat(psi, all_halo=.true.)

		! do j = mesh%full_lat_ibeg + 1, mesh%full_lat_iend - 1
		! 	do i = mesh%half_lon_ibeg, mesh%half_lon_iend
		! 		u(i,j,1) = (psi(i,j) - psi(i,j-1)) / mesh%le_lon(j) * radius
		! 	end do
		! end do
  !   call fill_zonal_halo_lon(u, all_halo=.true.)
		
		! do j = mesh%half_lat_ibeg, mesh%half_lat_iend
		! 	do i = mesh%full_lon_ibeg, mesh%full_lon_iend
		! 		v(i,j,1) = - (psi(i,j) - psi(i-1,j)) / mesh%le_lat(j) * radius
		! 	end do
		! end do
				

		do j = mesh%full_lat_ibeg + 1, mesh%full_lat_iend - 1
			lat = mesh%full_lat(j)
			do i = mesh%half_lon_ibeg, mesh%half_lon_iend
				lon = mesh%half_lon(i) - c1
				u(i,j,1) = k0 * sin(lon)**2 * sin(2 * lat) * cos_t + c2 * cos(lat)
			end do
		end do

		call fill_zonal_halo_lon(u, all_halo=.true.)

		do j = mesh%half_lat_ibeg, mesh%half_lat_iend
			lat = mesh%half_lat(j)
			do i = mesh%full_lon_ibeg, mesh%full_lon_iend
				lon = mesh%full_lon(i) - c1
				v(i,j,1) = k0 * sin(2 * lon) * cos(lat) * cos_t
			end do
		end do

		end associate

	end subroutine deform_case4_test_set_uv

end module deform_test_mod
